#include "DependencyEncoder2OSibling.h"
#include "math.h"
#include "ObjectReader.h"
#include "DependencyInstance.h"
#include <ctime>

DependencyEncoder2OSibling::DependencyEncoder2OSibling(int maxLength, DependencyPipe *pipe, char *trainforest, Parameters *param) : DependencyEncoder(maxLength, pipe, trainforest, param) {
	probs_trips = vector<vector<vector<double> > >(maxLength);
	probs_sibs = vector<vector<vector<double> > >(maxLength);
	for(int i = 0; i < maxLength; i++){
		probs_trips[i] = vector<vector<double> >(maxLength);
		probs_sibs[i] =  vector<vector<double> >(maxLength);
		for(int j = 0; j < maxLength; j++){
			int nn;
			if(j < i){
				nn = j;
			}
			else{
				nn = maxLength;
			}
			probs_trips[i][j] = vector<double>(nn);
			probs_sibs[i][j] = vector<double>(2);
		}
	}
}

DependencyEncoder2OSibling::~DependencyEncoder2OSibling(){}

int DependencyEncoder2OSibling::getKey(int s, int t, int dir, int comp, int length){
	int key = s;
	key = key * length + t;
	key = key * 2 + dir;
	key = key * 3 + comp;
	return key;
}

double DependencyEncoder2OSibling::getZ(int length, double *beta){
	int key, key1, key2;

	for(int i = 0; i < length; i++){
		key = getKey(i, i, 0, 1, length);
		beta[key] = 0.0;
		key = getKey(i, i, 1, 1, length);
		beta[key] = 0.0;
	}

	for(int j = 1; j < length; j++){
		for(int s = 0; s + j < length; s++){
			int t = s + j;
			double prodProb_st = probs[s][t][0];
			double prodProb_ts = probs[s][t][1];

			//init beta
			//incomplete spans
			//r == s
			int key_st_0 = getKey(s, t, 0, 0, length);
			double prodProb_sst = probs_trips[s][s][t] + probs_sibs[s][t][0] + prodProb_st;
			key1 = getKey(s, s, 0, 1, length);
			key2 = getKey(s + 1, t, 1, 1, length);
			beta[key_st_0] = logsumexp(beta[key_st_0], beta[key1] + beta[key2] + prodProb_sst, true);

			//r == t
			int key_ts_0 = getKey(s, t, 1, 0, length);
			double prodProb_tts = probs_trips[t][t][s] + probs_sibs[t][s][0] + prodProb_ts;
			key1 = getKey(s, t - 1, 0, 1, length);
			key2 = getKey(t, t, 1, 1, length);
			beta[key_ts_0] = logsumexp(beta[key_ts_0], beta[key1] + beta[key2] + prodProb_tts, true);

			//sibling spans
			int key_st_2 = getKey(s, t, 0, 2, length);
			beta[key_st_2] = 0.0;
			int key_ts_2 = getKey(s, t, 1, 2, length);
			beta[key_ts_2] = 0.0;
			bool flg_st_2 = true, flg_ts_2 = true;

			//complete spans
			int key_st_1 = getKey(s, t, 0, 1, length);
			beta[key_st_1] = 0.0;
			int key_ts_1 = getKey(s, t, 1, 1, length);
			beta[key_ts_1] = 0.0;
			bool flg_st_1 = true, flg_ts_1 = true;

			//calc sibling spans
			for(int r = s; r < t; r++){
				key1 = getKey(s, r, 0 ,1, length);
				key2 = getKey(r + 1, t, 1, 1, length);

				beta[key_st_2] = logsumexp(beta[key_st_2], beta[key1] + beta[key2], flg_st_2);
				flg_st_2 = false;

				beta[key_ts_2] = logsumexp(beta[key_ts_2], beta[key1] + beta[key2], flg_ts_2);
				flg_ts_2 = false;
			}
			
			//calc incomplete spans
			for(int r = s + 1; r < t; r++){
				key1 = getKey(s, r, 0, 0, length);
				key2 = getKey(r, t, 0, 2, length);
				double prodProb_srt = probs_trips[s][r][t] + probs_sibs[r][t][1] + prodProb_st;
				beta[key_st_0] = logsumexp(beta[key_st_0], beta[key1] + beta[key2] + prodProb_srt, false);

				key1 = getKey(s, r, 1, 2, length);
				key2 = getKey(r, t, 1, 0, length);
				double prodProb_trs = probs_trips[t][r][s] + probs_sibs[r][s][1] + prodProb_ts;
				beta[key_ts_0] = logsumexp(beta[key_ts_0], beta[key1] + beta[key2] + prodProb_trs, false);
			}

			//calc complete spans
			for(int r = s; r <= t; r++){
				if(r != s){
					key1 = getKey(s, r, 0, 0, length);
					key2 = getKey(r, t, 0, 1, length);
					beta[key_st_1] = logsumexp(beta[key_st_1], beta[key1] + beta[key2], flg_st_1);
					flg_st_1 = false;
				}
				if(r != t){
					key1 = getKey(s, r, 1, 1, length);
					key2 = getKey(r, t, 1, 0, length);
					beta[key_ts_1] = logsumexp(beta[key_ts_1], beta[key1] + beta[key2], flg_ts_1);
					flg_ts_1 = false;
				}
			}
		}
	}

	key1 = getKey(0, length - 1, 0, 1, length);
	key2 = getKey(0, length - 1, 1, 1, length);
	return logsumexp(beta[key1], beta[key2], false);
}

void DependencyEncoder2OSibling::inside_outside(int length, double *beta, double *alpha){
	int key;
	int end = length - 1;
	for(int d = 0; d < 2; d++){
		for(int c = 0 ; c < 3; c++){
			key = getKey(0, end, d, c, length);
			alpha[key] = 0.0;
		}
	}

	for(int j = end; j >= 1; j--){
		for(int s = 0; s + j < length; s++){
			int t = s + j;

			int key_a, key_b;

			//init alpha
			//sibling spans
			int key_st_2 = getKey(s, t, 0, 2, length);
			alpha[key_st_2] = 0.0;
			bool flg_st_2 = true;
			for(int r = 0; r < s; r++){
				double prodProb_rst = probs_trips[r][s][t] + probs_sibs[s][t][1] + probs[r][t][0];
				key_b = getKey(r, s, 0, 0, length);
				key_a = getKey(r, t, 0, 0, length);
				alpha[key_st_2] = logsumexp(alpha[key_st_2], beta[key_b] + alpha[key_a] + prodProb_rst, flg_st_2);
				flg_st_2 = false;
			}
			for(int r = t + 1; r < length; r++){
				double prodProb_rts = probs_trips[r][t][s] + probs_sibs[t][s][1] + probs[s][r][1];
				key_b = getKey(t, r, 1, 0, length);
				key_a = getKey(s, r, 1, 0, length);
				alpha[key_st_2] = logsumexp(alpha[key_st_2], beta[key_b] + alpha[key_a] + prodProb_rts, flg_st_2);
				flg_st_2 = false;
			}

			//complete spnas
			int key_st_1 = getKey(s, t, 0, 1, length);
			bool flg_st_1 = true;
			alpha[key_st_1] = 0.0;
			if(t + 1 < length){
				key_a = getKey(s, t + 1, 1, 0, length);
				double prodProb = probs_trips[t + 1][t + 1][s] + probs_sibs[t + 1][s][0] + probs[s][t + 1][1];
				alpha[key_st_1] = logsumexp(alpha[key_st_1], alpha[key_a] + prodProb, flg_st_1);
				flg_st_1 = false;
			}

			int key_ts_1 = getKey(s, t, 1, 1, length);
			bool flg_ts_1 = true;
			alpha[key_ts_1] = 0.0;
			if(s != 0){
				key_a = getKey(s - 1, t, 0, 0, length);
				double prodProb = probs_trips[s - 1][s - 1][t] + probs_sibs[s - 1][t][0] + probs[s - 1][t][0];
				alpha[key_ts_1] = logsumexp(alpha[key_ts_1], alpha[key_a] + prodProb, flg_ts_1);
				flg_ts_1 = false;
			}

			for(int r = 0; r < s; r++){
				key_b = getKey(r, s, 0, 0, length);
				key_a = getKey(r, t, 0, 1, length);
				alpha[key_st_1] = logsumexp(alpha[key_st_1], beta[key_b] + alpha[key_a], flg_st_1);
				flg_st_1 = false;

				if(!((r == 0) && (t == length -1))){
					key_b = getKey(r, s - 1, 0 ,1, length);
					key_a = getKey(r, t, 0, 2, length);
					alpha[key_ts_1] = logsumexp(alpha[key_ts_1], beta[key_b] + alpha[key_a], flg_ts_1);
					flg_ts_1 = false;
				}
			}
			for(int r = t + 1; r < length; r++){
				if(!((s == 0) && (r == length -1))){
					key_b = getKey(t + 1, r, 1, 1, length);
					key_a = getKey(s, r, 0 ,2, length);
					alpha[key_st_1] = logsumexp(alpha[key_st_1], beta[key_b] + alpha[key_a], flg_st_1);
					flg_st_1 = false;
				}

				key_b = getKey(t, r, 1, 0, length);
				key_a = getKey(s, r, 1, 1, length);
				alpha[key_ts_1] = logsumexp(alpha[key_ts_1], beta[key_b] + alpha[key_a], flg_ts_1);
				flg_ts_1 = false;
			}

			//incomplete spans
			int key_st_0 = getKey(s, t, 0, 0, length);
			alpha[key_st_0] = 0.0;
			bool flg_st_0 = true;

			int key_ts_0 = getKey(s, t, 1, 0, length);
			alpha[key_ts_0] = 0.0;
			bool flg_ts_0 = true;

			for(int r = t; r < length; r++){
				key_b = getKey(t, r, 0 ,1, length);
				key_a = getKey(s, r, 0 ,1, length);
				alpha[key_st_0] = logsumexp(alpha[key_st_0], beta[key_b] + alpha[key_a], flg_st_0);
				flg_st_0 = false;

				if(r != t){
					key_b = getKey(t, r, 0, 2, length);
					key_a = getKey(s, r, 0, 0, length);
					double prodProb_str = probs_trips[s][t][r] + probs_sibs[t][r][1] + probs[s][r][0];
					alpha[key_st_0] = logsumexp(alpha[key_st_0], beta[key_b] + alpha[key_a] + prodProb_str, flg_st_0);
					flg_st_0 = false;
				}
			}

			for(int r = 0; r <= s; r++){
				key_b = getKey(r, s, 1, 1, length);
				key_a = getKey(r, t, 1, 1, length);
				alpha[key_ts_0] = logsumexp(alpha[key_ts_0], beta[key_b] + alpha[key_a], flg_ts_0);
				flg_ts_0 = false;

				if(r != s){
					key_b = getKey(r, s, 0, 2, length);
					key_a = getKey(r, t, 1, 0, length);
					double prodProb_tsr = probs_trips[t][s][r] + probs_sibs[s][r][1] + probs[r][t][1];
					alpha[key_ts_0] = logsumexp(alpha[key_ts_0], beta[key_b] + alpha[key_a] + prodProb_tsr, flg_ts_0);
					flg_ts_0 = false;
				}
			}
		}
	}
}

void DependencyEncoder2OSibling::getExpected(long double *expected, double *beta, double *alpha, double z, int length){
	// Get production crap.
	for (int w1 = 0; w1 < length; w1++) {
		for (int w2 = w1 + 1; w2 < length; w2++) {
			for (int ph = 0; ph < 2; ph++) {
				int *keys = NULL;
				int size = in->readArray(keys);
				double m = 0.0;
				if(ph == 0){
					//w1 = s, w2 = t
					int key_b = getKey(w1 + 1, w2, 1, 1, length);
					int key_a = getKey(w1, w2, 0, 0, length);
					double prodProb_sst = probs_trips[w1][w1][w2] + probs_sibs[w1][w2][0] + probs[w1][w2][0];
					m += exp(beta[key_b] + alpha[key_a] + prodProb_sst -z);
					for(int r = w1 + 1; r < w2; r++){
						//w1 = s, r = r, w2 = t;
						key_b = getKey(w1, r, 0, 0, length);
						int key_b2 = getKey(r, w2, 0, 2, length);
						key_a = getKey(w1, w2, 0, 0, length);
						double prodProb_srt = probs_trips[w1][r][w2] + probs_sibs[r][w2][1] + probs[w1][w2][0];
						m += exp(beta[key_b] + beta[key_b2] + alpha[key_a] + prodProb_srt - z);
					}
				}
				else{
					//w1 = s, w2 = t
					int key_b = getKey(w1, w2 - 1, 0, 1, length);
					int key_a = getKey(w1, w2, 1, 0, length);
					double prodProb_tts = probs_trips[w2][w2][w1] + probs_sibs[w2][w1][0] + probs[w1][w2][1];
					m += exp(beta[key_b] + alpha[key_a] + prodProb_tts - z);
					for(int r = w1 + 1; r < w2; r++){
						//w1 = s, r = r, w2 = t;
						key_b = getKey(w1, r, 0, 2, length);
						int key_b2 = getKey(r, w2, 1, 0, length);
						key_a = getKey(w1, w2, 1, 0, length);
						double prodProb_trs = probs_trips[w2][r][w1] + probs_sibs[r][w1][1] + probs[w1][w2][1];
						m += exp(beta[key_b] + beta[key_b2] + alpha[key_a] + prodProb_trs - z);
					}
				}
				updataExpected(expected, keys, m, size);
				delete[](keys);
			}
		}
	}
	int last = in->readInt();
	if(last != -3){
		cout<<"Error reading file."<<endl;
		exit(0);
	}
	
	for(int w1 = 0; w1 < length; w1++){
		for(int w2 = w1; w2 < length; w2++){
			for(int w3 = w2 + 1; w3 < length; w3++){
				int *keys = NULL;
				int size = in->readArray(keys);
				double m = 0.0;
				if(w1 == w2){
					//w1 = w2 = s, w3 = t
					int key_b = getKey(w1 + 1, w3, 1, 1, length);
					int key_a = getKey(w1, w3, 0, 0, length);
					double prodProb_sst = probs_trips[w1][w2][w3] + probs_sibs[w2][w3][0] + probs[w1][w3][0];
					m += exp(beta[key_b] + alpha[key_a] + prodProb_sst -z);
				}
				else{
					//w1 = s, w2 = r, w3 = t
					int key_b = getKey(w1, w2, 0, 0, length);
					int key_b2 = getKey(w2, w3, 0, 2, length);
					int key_a = getKey(w1, w3, 0, 0, length);
					double prodProb_srt = probs_trips[w1][w2][w3] + probs_sibs[w2][w3][1] + probs[w1][w3][0];
					m += exp(beta[key_b] + beta[key_b2] + alpha[key_a] + prodProb_srt - z);
				}
				updataExpected(expected, keys, m, size);
				delete[](keys);
			}
		}
		for(int w2 = w1; w2 >= 0; w2--){
			for(int w3 = w2 - 1; w3 >= 0; w3--){
				int *keys = NULL;
				int size = in->readArray(keys);
				double m = 0.0;
				if(w1 == w2){
					//w3 = s, w2 = w1 = t
					int key_b = getKey(w3, w1 - 1, 0, 1, length);
					int key_a = getKey(w3, w1, 1, 0, length);
					double prodProb_tts = probs_trips[w1][w2][w3] + probs_sibs[w2][w3][0] + probs[w3][w1][1];
					m += exp(beta[key_b] + alpha[key_a] + prodProb_tts - z);
				}
				else{
					//w3 = s, w2 = r, w1 = t
					int key_b = getKey(w3, w2, 0, 2, length);
					int key_b2 = getKey(w2, w1, 1, 0, length);
					int key_a = getKey(w3, w1, 1, 0, length);
					double prodProb_trs = probs_trips[w1][w2][w3] + probs_sibs[w2][w3][1] + probs[w3][w1][1];
					m += exp(beta[key_b] + beta[key_b2] + alpha[key_a] + prodProb_trs - z);
				}
				updataExpected(expected, keys, m, size);
				delete[](keys);
			}
		}
	}
	last = in->readInt();
	if(last != -3){
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	for(int w1 = 0; w1 < length; w1++){
		for(int w2 = 0; w2 <length; w2++){
			for(int wh = 0; wh < 2; wh++){
				if(w1 != w2){
					int *keys = NULL;
					int size = in->readArray(keys);
					double m = 0.0;
					if(wh == 0){
						if(w1 < w2){
							//w1 = s, w2 = t;
							int key_b = getKey(w1 + 1, w2, 1, 1, length);
							int key_a = getKey(w1, w2, 0, 0, length);
							double prodProb_sst = probs_trips[w1][w1][w2] + probs_sibs[w1][w2][0] + probs[w1][w2][0];
							m += exp(beta[key_b] + alpha[key_a] + prodProb_sst -z);
						}
						else{
							//w1 = t, w2 = s
							int key_b = getKey(w2, w1 - 1, 0, 1, length);
							int key_a = getKey(w2, w1, 1, 0, length);
							double prodProb_tts = probs_trips[w1][w1][w2] + probs_sibs[w1][w2][0] + probs[w2][w1][1];
							m += exp(beta[key_b] + alpha[key_a] + prodProb_tts - z);
						}
					}
					else{
						if(w1 < w2){
							//s = s, w1 = r, w2 = t
							for(int s = 0; s < w1; s++){
								int key_b = getKey(s, w1, 0, 0, length);
								int key_b2 = getKey(w1, w2, 0, 2, length);
								int key_a = getKey(s, w2, 0, 0, length);
								double prodProb_srt = probs_trips[s][w1][w2] + probs_sibs[w1][w2][1] + probs[s][w2][0];
								m += exp(beta[key_b] + beta[key_b2] + alpha[key_a] + prodProb_srt - z);
							}
						}
						else{
							//w2 = s, w1 = r, t = t
							for(int t = w1 + 1; t < length; t++){
								int key_b = getKey(w2, w1, 0, 2, length);
								int key_b2 = getKey(w1, t, 1, 0, length);
								int key_a = getKey(w2, t, 1, 0, length);
								double prodProb_trs = probs_trips[t][w1][w2] + probs_sibs[w1][w2][1] + probs[w2][t][1];
								m += exp(beta[key_b] + beta[key_b2] + alpha[key_a] + prodProb_trs - z);
							}
						}
					}
					updataExpected(expected, keys, m, size);
					delete[](keys);
				}
			}
		}
	}
	last = in->readInt();
	if(last != -3){
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	int *keys = NULL;
	int size = in->readArray(keys);
	for(int i = 0; i < size; i++){
		int j = keys[i];
		expected[j] -= 1.0;
	}
	
	delete[](keys);
	last = in->readInt();
	if (last != -4) {
		cout<<"Error reading file."<<endl;
		exit(0);
	}

	DependencyInstance* inst = new DependencyInstance();
	inst->readObject(*in);

	last = in->readInt();
	if (last != -1) {
		cout<<"Error reading file."<<endl;
		exit(0);
	}
	delete(inst);
}

double DependencyEncoder2OSibling::gradient(long double *expected, int length){
	long offset = in->objTell();
	DependencyInstance *inst = ((DependencyPipe2OSibling*)pipe)->readInstance(*in, length, probs, probs_trips, probs_sibs, nt_probs, param);
	if(0 != (in->objSeek(offset))){
		cout<<"error on fseek of trainforset!"<<endl;
		exit(0);
	}

	double obj = 0.0;
	double *beta = new double[length * length * 2 * 3];
	double *alpha = new double[length * length * 2 * 3];

	double z = getZ(length, beta);

	inside_outside(length, beta, alpha);

	obj = z - param->getScore(inst->fv);
	delete(inst);

	getExpected(expected, beta, alpha, z, length);

	delete[](beta);
	delete[](alpha);

	return obj;
}
